#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn
# @Date  : 2024-10-04

"""
【节点名称】：
    SpireFormatReaderNode
【依赖项安装】：
    pip install spirems
【订阅类型】：
    无
【发布类型】：
    sensor_msgs::CompressedImage （输出图像）
【构造参数说明】：
    parameter_file (str): 全局参数文件
【节点参数】：
    resize (list): 图像重采样分辨率
    fps (int): 输出帧率
【备注】：
    无
"""
import os
import threading
import time
import json
import cv2
import uuid
from typing import Union
from queue import Queue
from spirems import Publisher, Subscriber, cvimg2sms, def_msg, QoS
from spirecv.base.BaseNode import BaseNode


def load_spire_dir(spire_dir, with_relative_path=False):
    """
    :param spire_dir: path to spire dataset root dir
    :param with_relative_path: output relative path or not
    :return: List - img_list, List - ann_list
    """
    img_list = []
    ann_list = []
    rel_img_list = []
    rel_ann_list = []
    for sub_dir in os.listdir(spire_dir):
        sub_fn = os.path.join(spire_dir, sub_dir)
        if os.path.isdir(sub_fn):
            if sub_dir == 'scaled_images':
                img_list.extend([os.path.join(sub_fn, img_nm) for img_nm in os.listdir(sub_fn) if
                                 os.path.splitext(img_nm)[-1] in ['.jpg', '.JPG', '.png']])
                ann_list.extend([os.path.join(spire_dir, 'annotations', img_nm + '.json') for img_nm in
                                 os.listdir(sub_fn) if
                                 os.path.splitext(img_nm)[-1] in ['.jpg', '.JPG', '.png']])
                if with_relative_path:
                    rel_img_list.extend([os.path.join(sub_dir, img_nm) for img_nm in os.listdir(sub_fn) if
                                        os.path.splitext(img_nm)[-1] in ['.jpg', '.JPG', '.png']])
                    rel_ann_list.extend([os.path.join('annotations', img_nm + '.json') for img_nm in
                                        os.listdir(sub_fn) if
                                        os.path.splitext(img_nm)[-1] in ['.jpg', '.JPG', '.png']])
            else:
                for sub_sub_dir in os.listdir(sub_fn):
                    sub_sub_fn = os.path.join(sub_fn, sub_sub_dir)
                    if os.path.isdir(sub_sub_fn):
                        if sub_sub_dir == 'scaled_images':
                            img_list.extend([os.path.join(sub_sub_fn, img_nm) for img_nm in os.listdir(sub_sub_fn) if
                                             os.path.splitext(img_nm)[-1] in ['.jpg', '.JPG', '.png']])
                            ann_list.extend([os.path.join(sub_fn, 'annotations', img_nm + '.json') for img_nm in
                                             os.listdir(sub_sub_fn) if
                                             os.path.splitext(img_nm)[-1] in ['.jpg', '.JPG', '.png']])
                            if with_relative_path:
                                rel_img_list.extend([os.path.join(sub_dir, sub_sub_dir, img_nm)
                                                     for img_nm in os.listdir(sub_sub_fn)
                                                     if os.path.splitext(img_nm)[-1] in ['.jpg', '.JPG', '.png']])
                                rel_ann_list.extend([os.path.join(sub_dir, 'annotations', img_nm + '.json')
                                                     for img_nm in os.listdir(sub_sub_fn)
                                                     if os.path.splitext(img_nm)[-1] in ['.jpg', '.JPG', '.png']])
    if with_relative_path:
        img_list.sort()
        ann_list.sort()
        rel_img_list.sort()
        rel_ann_list.sort()
        return img_list, ann_list, rel_img_list, rel_ann_list
    else:
        img_list.sort()
        ann_list.sort()
        return img_list, ann_list


class SpireFormatReaderNode(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        sms_shutdown_emit: bool = True
    ):
        threading.Thread.__init__(self)
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file,
            sms_shutdown=True
        )
        self.sms_shutdown_emit = sms_shutdown_emit
        self.client_id = str(uuid.uuid4()).replace('-', '_')
        self.img_i_queue = Queue()
        self.queue_pool.append(self.img_i_queue)
        self.img_i = 0
        self._next_reader = Subscriber(
            '/' + job_name + '/launch_next', 'std_msgs::Boolean', self.launch_next,
            ip=ip, port=port, qos=QoS.Reliability
        )
        self._image_writer = Publisher(
            '/' + job_name + '/sensor/image_raw', 'sensor_msgs::CompressedImage',
            ip=ip, port=port, qos=QoS.Reliability
        )
        self._result_writer = Publisher(
            '/' + job_name + '/gt/results', 'spirecv_msgs::2DTargets',
            ip=ip, port=port, qos=QoS.Reliability
        )
        self.frame_id = 0
        self.spire_root = self.get_param("spire_root", "D:/dataset/GT")
        self.dataset_name = self.get_param("dataset_name", "coco_detection")
        self.categories = self.get_param("categories", ["boat"])
        self.next_auto = self.get_param("next_auto", 0)
        self.img_list, self.ann_list, self.rel_img_list, self.rel_ann_list = load_spire_dir(
            self.spire_root, with_relative_path=True)

        self.start()

    def release(self):
        BaseNode.release(self)
        self._image_writer.kill()
        self._next_reader.kill()
        self._result_writer.kill()

    def launch_next(self, msg: dict = None):
        if (isinstance(msg, dict) and msg['data']) or msg is None:
            if self.img_i < len(self.img_list):
                self.img_i_queue.put(self.img_i)
                self.img_i += 1
            else:
                print('Finish! {} / {}'.format(self.img_i, len(self.img_list)))

    def run(self):
        while self.is_running():
            img_i = self.img_i_queue.get(block=True)
            if img_i is None:
                break

            img = cv2.imread(self.img_list[img_i])
            rel_path = self.rel_img_list[img_i]
            parts = os.path.normpath(rel_path).split(os.sep)

            with open(self.ann_list[img_i], 'r') as f:
                gt = json.load(f)

            assert parts[-1] == gt['file_name']

            tgt_msg = def_msg('spirecv_msgs::2DTargets')
            tgt_msg['file_name'] = gt['file_name']
            tgt_msg['height'] = gt['height']
            tgt_msg['width'] = gt['width']
            tgt_msg['dataset'] = self.dataset_name
            tgt_msg['client_id'] = self.client_id
            tgt_msg['img_id'] = img_i
            tgt_msg['img_total'] = len(self.img_list)
            tgt_msg['targets'] = []
            for anno in gt['annos']:
                if anno['category_name'] in self.categories:
                    tgt_msg['targets'].append(anno)
            self._result_writer.publish(tgt_msg)
            print(tgt_msg)

            msg = cvimg2sms(img)
            msg['img_id'] = img_i
            msg['img_total'] = len(self.img_list)
            msg['file_name'] = os.path.basename(self.img_list[img_i])
            msg['client_id'] = self.client_id
            if len(parts) == 3:
                msg['video_name'] = parts[0]

            msg['spirecv_msgs::2DTargets'] = tgt_msg
            self._image_writer.publish(msg)

            if self.next_auto:
                time.sleep(1)
                self.launch_next()

        self.release()
        print('{} quit!'.format(self.__class__.__name__))


if __name__ == '__main__':
    parameter_file = '../../params/spirecv2/default_params.json'
    node = SpireFormatReaderNode('MOD', param_dict_or_file=parameter_file)
    node.launch_next()
